# 批量标注图片点位
import base64
import os
from traceback import format_exc

from PIL import Image, ImageDraw, ImageFont

from app.lib import yolo, baidu
from app.resource import resource_dir
from app.utils.index import orientation, PIL2base64


def mark(file_path, save_file):
  # 先加载图片进行大小检测和方向校正
  img_original = Image.open(file_path)
  if img_original.mode == "RGBA":
    img_original = img_original.convert("RGB")
  img_fix = orientation(img_original)

  bs64 = PIL2base64(img_fix)
  # 调用yolo接口
  top = yolo.get_photo_head_param(bs64)
  # 调用百度接口
  param = baidu.get_photo_mark_param(bs64)

  # 绘制
  line_color = (255, 0, 0)
  line_color2 = (200, 0, 0)
  rect_color = (240, 0, 0)
  line_width = 2

  font_file_path = os.path.join(resource_dir, "simhei.ttf")
  font = ImageFont.truetype(font_file_path, 20)  # 指定字体和大小
  font_color = (255, 69, 0)
  source_width, source_height = img_fix.size
  draw = ImageDraw.Draw(img_fix)

  hor_center = source_width / 2  # 横向中线
  ver_center = source_height / 2  # 竖向中线

  # 头顶位置
  draw.line([(hor_center - 100, top), (hor_center + 100, top)], fill=line_color, width=line_width)
  # 脸顶位置
  draw.line([(hor_center - 100, param.options.face_top.y), (hor_center + 100, param.options.face_top.y)], fill=line_color, width=line_width)
  # 文字注释
  draw.line([(hor_center, top), (hor_center, param.options.face_top.y)], fill=line_color2, width=line_width)
  draw.text((hor_center, top+(param.options.face_top.y-top)/2), f"长度{(param.options.face_top.y-top):.2f}", fill=font_color, font=font)
  # 人脸框
  draw.line([(param.options.face_location.left, param.options.face_location.top), (param.options.face_location.left+param.options.face_location.width, param.options.face_location.top)], fill=line_color, width=line_width)
  draw.line([(param.options.face_location.left, param.options.face_location.top+param.options.face_location.height), (param.options.face_location.left+param.options.face_location.width, param.options.face_location.top+param.options.face_location.height)], fill=line_color, width=line_width)
  draw.line([(param.options.face_location.left, param.options.face_location.top), (param.options.face_location.left, param.options.face_location.top+param.options.face_location.height)], fill=line_color, width=line_width)
  draw.line([(param.options.face_location.left+param.options.face_location.width, param.options.face_location.top), (param.options.face_location.left+param.options.face_location.width, param.options.face_location.top+param.options.face_location.height)], fill=line_color, width=line_width)
  # 文字注释
  draw.text((param.options.face_location.left+param.options.face_location.width, param.options.face_location.top+param.options.face_location.height/2), f"长度{param.options.face_location.height:.2f}", fill=font_color, font=font)
  draw.text((param.options.face_location.left+param.options.face_location.width/2, param.options.face_location.top+param.options.face_location.height), f"宽度{param.options.face_location.width:.2f}", fill=font_color, font=font)
  # 文字注释
  draw.line([(hor_center, param.options.face_top.y), (hor_center, param.options.face_location.top)], fill=line_color2, width=line_width)
  draw.text((hor_center, param.options.face_top.y+(param.options.face_location.top-param.options.face_top.y)/2), f"长度{(param.options.face_location.top-param.options.face_top.y):.2f}", fill=font_color, font=font)

  img_fix.save(save_file)


def mark_photo_one(file_path, save_dir):
  print("开始处理" + file_path)
  file_name = os.path.basename(file_path)
  base_name, file_ext = os.path.splitext(file_name)
  save_file = os.path.join(save_dir, file_name)
  error_file = os.path.join(save_dir, f"{base_name}_错误.txt")

  try:
    mark(file_path, save_file)
  except Exception as exc:
    with open(error_file, "w") as file:
      file.write(format_exc())


if __name__=="__main__":
  source_dir = r"E:\211\人脸"

  save_dir = source_dir + "_标注"
  if not os.path.exists(save_dir):
    os.mkdir(save_dir)

  for root, dirs, files in os.walk(source_dir):
    index = 1
    total = len(files)
    for file in files:
      print("--", index, "/", total, "--")
      file_path = os.path.join(root, file)
      mark_photo_one(file_path, save_dir)
      index += 1

  print("全部处理完毕")
